package graph2.weight.prim;

import graph.weight.prim2.IndexMinHeap;
import graph2.weight.Edge;
import graph2.weight.WeightedGraph;

import java.util.ArrayList;
import java.util.LinkedList;
import java.util.List;

/**
 * 有权图的最小生成树 原理是取得连接到每个节点权重最小的边
 */
public class PrimMstReview<Weight extends Number & Comparable> {

    /**
     * 指向顶点的权重最小的边
     */
    private Edge<Weight>[] edgeTo;

    /**
     * 最小生成树的边
     */
    private List<Edge> edges;

    /**
     * 已经有最小的边指向
     */
    private boolean[] marked;

    /**
     * 图的总权重
     */
    private double totalWeight;

    /**
     * 索引是节点 值是指向此索引的最短边
     */
    private IndexMinHeap<Weight> indexMinHeap;

    /**
     * 图的引用
     */
    private WeightedGraph graph;


    public PrimMstReview(WeightedGraph graph) {
        edgeTo = (Edge<Weight>[]) new Edge[graph.V()];
        edges = new ArrayList<>();
        marked = new boolean[graph.V()];
        indexMinHeap = new IndexMinHeap<>(graph.V());

        visit(0);
        while (!indexMinHeap.isEmpty()) {
            int index = indexMinHeap.extractMinIndex();
            edges.add(edgeTo[index]);
            visit(index);
        }

        totalWeight = edges.get(0).weight().doubleValue();
        for (int i = 1; i < edges.size(); i++) {
            totalWeight += edges.get(i).weight().doubleValue();
        }

    }

    private void visit(int i) {
        marked[i] = true;
        Iterable<Edge<Weight>> adj = graph.adj(i);
        for (Edge<Weight> edge : adj) {
            int other = edge.other(i);
            if (!marked[other]) {
                if (edgeTo[other] == null) {
                    indexMinHeap.insert(other, edge.weight());
                    edgeTo[other] = edge;
                } else if (edgeTo[other].compareTo(edge) > 0) {
                    indexMinHeap.change(other, edge.weight());
                    edgeTo[other] = edge;
                }
            }
        }
    }


    // 返回最小生成树的所有边
    List<Edge> mstEdges() {
        return edges;
    }

    // 返回最小生成树的权值
    Number result() {
        return totalWeight;
    }


}
